"""统一的彩色打印工具模块"""

import datetime
from colorama import Fore, Style, init
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
from typing import Any, Optional

# 初始化colorama
init(autoreset=True)

class PrettyPrinter:
    """统一的彩色打印器"""
    
    def __init__(self, show_timestamp: bool = True, show_icons: bool = True):
        self.show_timestamp = show_timestamp
        self.show_icons = show_icons
        
    def _get_message_style(self, message) -> tuple[str, str, str]:
        """获取消息的样式配置"""
        if isinstance(message, HumanMessage):
            return Fore.BLUE, "[用户]", "🧑" if self.show_icons else ""
        elif isinstance(message, AIMessage):
            return Fore.GREEN, "[AI助手]", "🤖" if self.show_icons else ""
        elif isinstance(message, ToolMessage):
            return Fore.YELLOW, "[工具]", "🔧" if self.show_icons else ""
        else:
            return Fore.WHITE, "[其他]", "ℹ️" if self.show_icons else ""
    
    def _format_timestamp(self) -> str:
        """格式化时间戳"""
        if not self.show_timestamp:
            return ""
        return f"[{datetime.datetime.now().strftime('%H:%M:%S')}] "
    
    def print_message(self, message: Any, role_type: Optional[str] = None):
        """打印消息"""
        timestamp = self._format_timestamp()
        
        if hasattr(message, "__class__") and hasattr(message, "content"):
            color, prefix, icon = self._get_message_style(message)
            icon_str = f"{icon} " if icon else ""
            
            if hasattr(message, "content") and message.content:
                content = str(message.content)
                # 处理长文本
                if len(content) > 200:
                    content = content[:200] + "..."
                print(f"{color}{timestamp}{icon_str}{prefix}: {content}{Style.RESET_ALL}")
            elif hasattr(message, "tool_calls") and message.tool_calls:
                print(f"{color}{timestamp}{icon_str}{prefix}: [工具调用] {message.tool_calls}{Style.RESET_ALL}")
            else:
                print(f"{color}{timestamp}{icon_str}{prefix}: {str(message)}{Style.RESET_ALL}")
        else:
            # 处理字符串或其他类型
            role_styles = {
                "human": (Fore.BLUE, "[用户]", "🧑"),
                "ai": (Fore.GREEN, "[AI助手]", "🤖"),
                "tool": (Fore.YELLOW, "[工具]", "🔧"),
                "system": (Fore.WHITE, "[系统]", "ℹ️")
            }
            
            color, prefix, icon = role_styles.get(role_type, role_styles["system"])
            icon_str = f"{icon} " if icon and self.show_icons else ""
            print(f"{color}{timestamp}{icon_str}{prefix}: {message}{Style.RESET_ALL}")
    
    def print_separator(self, title: str = "步骤", char: str = "=", width: int = 50):
        """打印分隔符"""
        separator = char * ((width - len(title) - 2) // 2)
        print(f"\n{Fore.CYAN}{separator} {title} {separator}{Style.RESET_ALL}\n")
    
    def print_header(self, title: str, char: str = "=", width: int = 60):
        """打印标题"""
        print(f"\n{Fore.MAGENTA}{char * width}{Style.RESET_ALL}")
        print(f"{Fore.MAGENTA}  {title}{Style.RESET_ALL}")
        print(f"{Fore.MAGENTA}{char * width}{Style.RESET_ALL}\n")
    
    def print_step_info(self, step_num: int, description: str):
        """打印步骤信息"""
        print(f"{Fore.CYAN}📋 步骤 {step_num}: {description}{Style.RESET_ALL}")
    
    def print_error(self, error_msg: str):
        """打印错误信息"""
        print(f"{Fore.RED}❌ 错误: {error_msg}{Style.RESET_ALL}")
    
    def print_success(self, success_msg: str):
        """打印成功信息"""
        print(f"{Fore.GREEN}✅ 成功: {success_msg}{Style.RESET_ALL}")
    
    def print_warning(self, warning_msg: str):
        """打印警告信息"""
        print(f"{Fore.YELLOW}⚠️ 警告: {warning_msg}{Style.RESET_ALL}")

# 创建全局实例
printer = PrettyPrinter()

# 便捷函数
def print_message(message, role_type=None):
    printer.print_message(message, role_type)

def print_separator(title="步骤"):
    printer.print_separator(title)

def print_header(title):
    printer.print_header(title)

def print_step_info(step_num, description):
    printer.print_step_info(step_num, description)

def print_error(error_msg):
    printer.print_error(error_msg)

def print_success(success_msg):
    printer.print_success(success_msg)

def print_warning(warning_msg):
    printer.print_warning(warning_msg)